package cn.edu.neu.softlab633.influencemaximization.sy.bean.CELF;

import cn.edu.neu.softlab633.influencemaximization.sy.bean.Graph;
import cn.edu.neu.softlab633.influencemaximization.sy.bean.MarginGain;
import cn.edu.neu.softlab633.influencemaximization.sy.model.lt.LinearThresholdModel;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.TreeMap;

/**
 * Created by Jason on 2017/5/11.
 */
public class CELFQueue {
    private TreeMap<CELFQueueEle, Integer> celfQueue;
    private TreeMap<Integer, CELFQueueEle> reverseCelfQueue;

    public CELFQueue() {
        celfQueue = new TreeMap<>();
        reverseCelfQueue = new TreeMap<>();
    }

    public void initCELFQueue(double threshold, Double[] query, Graph graph) {
        if (celfQueue == null || reverseCelfQueue == null) {
            celfQueue = new TreeMap<>();
            reverseCelfQueue = new TreeMap<>();
        }
        Iterator<Integer> iterator = graph.getOutgoingIndex().getIndex().keySet().iterator();
        while (iterator.hasNext()) {
            Integer next = iterator.next();
            ArrayList<Integer> tmp = new ArrayList<>();
            tmp.add(next);
            MarginGain gain = LinearThresholdModel.influencePropagation(threshold, query, tmp, graph);
            this.add(new CELFQueueEle(next, gain, 1));
        }
    }

    public CELFQueue(TreeMap<CELFQueueEle, Integer> celfQueue) {
        this.celfQueue = celfQueue;
    }

    public CELFQueueEle getFirst() {
        return celfQueue.firstKey();
    }

    public CELFQueueEle getLast() {
        return celfQueue.lastKey();
    }

    public Integer add(CELFQueueEle gain) {
        // 修改代码
        reverseCelfQueue.put(gain.getId(), gain);
        return celfQueue.put(gain, gain.getId());
    }

    public boolean removeAll(Collection<Integer> gain) {
        Iterator<Integer> iterator = gain.iterator();
        while (iterator.hasNext()) {
            int integer = iterator.next();
            if (reverseCelfQueue.containsKey(integer)) {
                celfQueue.remove(reverseCelfQueue.get(integer));
                reverseCelfQueue.remove(integer);
            }
        }
        return true;
    }

    public CELFQueueEle peek() {
        CELFQueueEle tmp = this.getFirst();
        celfQueue.remove(tmp);
        reverseCelfQueue.remove(tmp.getId());
        return tmp;
    }

    public int size() {
        return celfQueue.size();
    }
}
